package com.sakura.前缀树;

import java.util.Arrays;

public class Code01_Trie {

    // 类前缀树
    class Trie {
        TreeNode root;

        class TreeNode {
            int pass;
            int end;
            TreeNode[] nexts;
            public TreeNode() {
                pass = 0;
                end = 0;
                nexts = new TreeNode[26];
            }
        }

        public Trie() {
            root = new TreeNode();
        }

        public void insert(String word) {
            TreeNode node = root;
            node.pass++;
            for (int i = 0, path; i < word.length(); i++) {
                path = word.charAt(i) - 'a';
                if (node.nexts[path] == null) {
                    node.nexts[path] = new TreeNode();
                }
                node = node.nexts[path];
                node.pass++;
            }
            node.end++;
        }

        public int search(String word) {
            TreeNode node = root;
            for (int i = 0, path; i < word.length(); i++) {
                path = word.charAt(i) - 'a';
                if (node.nexts[path] == null) {
                    return 0; // word单词从来没有加入到前缀树中
                }
                node = node.nexts[path];
            }
            return node.end;
        }

        public int prefixNumber(String prefix) {
            TreeNode node = root;
            for (int i = 0, path; i < prefix.length(); i++) {
                path = prefix.charAt(i) - 'a';
                if (node.nexts[path] == null) {
                    return 0; // word单词从来没有加入到前缀树中
                }
                node = node.nexts[path];
            }
            return node.pass;
        }

        public void erase(String word) {
            if (search(word) > 0) {
                TreeNode node = root;
                node.pass--;
                for (int i = 0, path; i < word.length(); i++) {
                    path = word.charAt(i) - 'a';
                    if (--node.pass == 0) {
                        node.nexts[path] = null;
                        return;
                    }
                    node = node.nexts[path];
                }
                node.end--;
            }
        }
    }

    // 静态数组实现
    class Trie2 {

        private static final int MAX = 150001;
        private static int[][] trie = new int[MAX][26];
        private static int[] pass = new int[MAX];
        private static int[] end = new int[MAX];
        private static int cnt;
        public static void build() {
            cnt = 1;
        }
        public static void insert(String word) {
            int cur = 1; // 类中的root
            pass[cur]++;
            for (int i = 0, path; i < word.length(); i++) {
                path = word.charAt(i) - 'a';
                if (trie[cur][path] == 0) {
                    trie[cur][path] = ++cnt;
                }
                cur = trie[cur][path];
                pass[cur]++;
            }
            end[cur]++;
        }

        public static int search(String word) {
            int cur = 1;
            for (int i = 0, path; i < word.length(); i++) {
                path = word.charAt(i) - 'a';
                if (trie[cur][path] == 0) {
                    return 0;
                }
                cur = trie[cur][path];
            }
            return end[cur];
        }

        public static int prefixNumber(String prefix) {
            int cur = 1;
            for (int i = 0, path; i < prefix.length(); i++) {
                path = prefix.charAt(i) - 'a';
                if (trie[cur][path] == 0) {
                    return 0;
                }
                cur = trie[cur][path];
            }
            return pass[cur];
        }

        public static void erase(String word) {
            if (search(word) > 0) {
                int cur = 1;
                pass[cur]--;
                for (int i = 0, path; i < word.length(); i++) {
                    path = word.charAt(i) - 'a';
                    if (--pass[cur] == 0) {
                        trie[cur][path] = 0;
                        return;
                    }
                    cur = trie[cur][path];
                    pass[cur]--;
                }
                end[cur]--;
            }
        }

        public static void clear() {
            for (int i = 1; i <= cnt; i++) {
                Arrays.fill(trie[i], 0);
                pass[i] = 0;
                end[i] = 0;
            }
            cnt = 1;
        }
    }
}
